feat(guardrails): add Airia guardrail integration - #40784

TLDR

Problem this solves:

• Airia customers on LiteLLM have no way to enforce Airia guardrails on proxy traffic
• Today they must copy a Python file onto the proxy host and reference it by dotted path

How it solves it:

• Adds airia as a built-in guardrail provider, configured with api_base + api_key
• Prompts are checked before the model call (pre_call), responses after it (post_call)
• Airia returns allow / block / redact; redactions are applied before the text moves on
• Streamed responses are moderated whole, then delivered redacted (incremental_diff + streaming_end_of_stream_only)
• A Gateway response that is not a JSON object, or a rewrite whose item count does not match what was sent, fails closed instead of erroring or misaligning the positional write-back

User Flow

Before: a proxy admin who wants Airia guardrails cannot reference the provider by name

1. They add guardrail: airia under guardrails: in config.yaml and start the proxy
2. Startup logs Skipping guardrail 'airia': invalid configuration ... Unsupported guardrail: airia
3. Every POST to http://litellm-domain/v1/chat/completions reaches the model unchecked

After: the same config enforces the Airia policy on every call

1. They add guardrail: airia under guardrails: in config.yaml and start the proxy
2. The proxy starts with the guardrail active
3. POST /v1/chat/completions with "my aws key is AKIAIOSFODNN7EXAMPLE" returns 400 — Guardrail raised an exception, Guardrail: airia, Message: <the message configured on their Airia guardrail>. The model is never called
4. POST with "Contact me at ada@example.com" returns 200; the model received Contact me at [EmailAddress1]
5. POST with a benign prompt returns 200 unchanged

Relevant issues

—

Pre-Submission checklist

• I have added meaningful tests
• The handful of test files covering my change pass locally: uv run pytest tests/test_litellm/proxy/guardrails/guardrail_hooks/test_airia.py -v → 21 passed
• My PR passes all required CI/CD checks
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• Greptile Confidence Score of at least 4/5

Screenshots / Proof of Fix

Setup shared by both runs. config.yaml:

model_list:
  - model_name: gpt-4o-mini
    litellm_params:
      model: openai/harness-echo          # local echo endpoint that records exactly what it received,
      api_base: http://127.0.0.1:8931     # so the redacted text the MODEL saw is observable
      api_key: not-a-real-key
guardrails:
  - guardrail_name: "airia"
    litellm_params:
      guardrail: airia
      mode: ["pre_call", "post_call"]
      default_on: true
      api_base: os.environ/AIRIA_GATEWAY_URL   # a running Airia AI Gateway; policy: PII → redact, Secrets → block
      api_key: os.environ/AIRIA_API_KEY
general_settings:
  master_key: sk-1234

The guardrail side is a live Airia deployment, not a mock. The model is a local echo endpoint because that is the only way to show what the provider actually received; nothing in the guardrail path is stubbed.

Before (a426dc4)

1. litellm --config config.yaml --port 4096
2. Output:

   LiteLLM Proxy:ERROR: init_guardrails.py:37 - Skipping guardrail 'airia': invalid configuration, proxy is starting WITHOUT this guardrail: Unsupported guardrail: airia
   

After (899e719, unit-tested; see Caveats for what changed since the run below)

(rebased onto the current base after the run below; the tree is byte-identical to the commit it was captured on)

Blocked — secret in the prompt

1. curl -s http://127.0.0.1:4097/v1/chat/completions -H 'Content-Type: application/json' -H 'Authorization: Bearer sk-1234' \
     -d '{"model":"gpt-4o-mini","messages":[{"role":"user","content":"my aws key is AKIAIOSFODNN7EXAMPLE"}]}'

2. HTTP 400

   {"error": {"message": "Guardrail raised an exception, Guardrail: airia, Message: Secret blocked by A51-317 demo policy.", "type": "None", "param": "None", "code": "400"}}

3. Echo endpoint call log: 0 calls for this request — the model was never invoked

Redacted — PII in the prompt

1. curl -s http://127.0.0.1:4097/v1/chat/completions -H 'Content-Type: application/json' -H 'Authorization: Bearer sk-1234' \
     -d '{"model":"gpt-4o-mini","messages":[{"role":"user","content":"Contact me at ada@example.com"}]}'

2. HTTP 200, choices[0].message.content = "echo: Contact me at [EmailAddress1]"
3. Echo endpoint call log shows the model received ['Contact me at [EmailAddress1]'] — the address never reached it

Allowed — benign prompt

1. curl -s http://127.0.0.1:4097/v1/chat/completions -H 'Content-Type: application/json' -H 'Authorization: Bearer sk-1234' \
     -d '{"model":"gpt-4o-mini","messages":[{"role":"user","content":"Explain recursion in one sentence."}]}'

2. HTTP 200, choices[0].message.content = "echo: Explain recursion in one sentence."
3. Echo endpoint call log shows the model received the prompt unchanged

Redacted — PII in the model's response (post_call)

The echo endpoint answers the fixed prompt please share the demo contact (no PII on the way in) with Sure, reach me at ada@example.com or call 415-555-0199.

1. curl -s http://127.0.0.1:4097/v1/chat/completions -H 'Content-Type: application/json' -H 'Authorization: Bearer sk-1234' \
     -d '{"model":"gpt-4o-mini","messages":[{"role":"user","content":"please share the demo contact"}]}'

2. HTTP 200, choices[0].message.content = "echo: Sure, reach me at [EmailAddress1] or call 415-555-0199."

Redacted — PII in a streamed response (post_call, stream: true)

1. curl -sN http://127.0.0.1:4097/v1/chat/completions -H 'Content-Type: application/json' -H 'Authorization: Bearer sk-1234' \
     -d '{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"please share the demo contact"}]}'

2. HTTP 200; the stream is held until the response is complete and Airia has answered, then one delta carries echo: Sure, reach me at [EmailAddress1] or call 415-555-0199.
3. grep -c 'ada@example.com' over the raw SSE bytes: 0

Type

🆕 New Feature

Caveats (if any)

Medium

• Streamed tool-call chunks reach the client before the end-of-stream verdict. _run_incremental_transform_stream in unified_guardrail.py yields a passthrough chunk for each tool-call delta as it arrives and only inspects the assembled tool calls against the guardrail's verdict once the stream ends, so a blocked tool call can be executed by the client before the terminal error lands. This is shared framework behavior, not specific to Airia: every guardrail using incremental_diff streaming has the same exposure. I looked for a way to close it from this guardrail alone and could not find one that does not regress the redaction path: streaming_buffer_until_moderated (which withholds every chunk until the full response is moderated) has no effect under incremental_diff — the mode branches away and returns before that flag is ever read — and switching to block_only to make it apply would silently drop REDACT verdicts on streamed text, replaying the original unredacted content on release. Closing this needs a change to the shared streaming framework (buffering tool-call chunks specifically, independent of text transform mode), which is out of scope for an isolated single-guardrail PR
• Separately, and outside streaming too: OpenAIChatCompletionsHandler.process_output_response skips every post_call guardrail entirely when a response has no text content (if not self._has_text_content(response): ... return response), so a tool-call-only response with no text bypasses Airia's response-side check regardless of streaming. Also shared framework behavior; also out of scope here

Low

• Streamed text arrives in one delta at end of stream rather than incrementally: the trade for redactions that can never underflow. Blocks still terminate the stream immediately

• The HTTP client is dependency-injected (async_handler constructor param) rather than mutated after construction, so tests build a fake client instead of monkeypatching an instance method

• during_call is deliberately not offered; it runs concurrently with the model call, so a block could land after the prompt reached the provider

• The proof's model is a local echo endpoint; the guardrail path is a live Airia deployment. No provider key is available to this machine for a paid-model run

• If Airia is unreachable the request is refused (blocked_content=False), not passed through — same fail-closed stance as the other guardrail providers

• Docs page will follow in a separate PR to litellm-docs

• BerriAI/litellm retired litellm_internal_staging as the trunk while this PR was open (main is now the sole default branch); rebased onto current main with no conflicts, and re-verified against the strict-ruleset rules that were added upstream in the same window (test return-type annotations, banned typing.cast, **kwargs typing). No behavior change, no budget-JSON edits

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR